from itertools import combinations

# 创建一个包含交易的示例数据集
dataset = [
    {'milk', 'bread', 'sugar'},
    {'milk', 'bread', 'butter'},
    {'milk', 'bread'},
    {'milk', 'sugar'},
    {'bread', 'butter'},
    {'bread', 'sugar'},
    {'milk', 'bread', 'butter'},
    {'milk', 'bread'},
    {'milk', 'butter'},
    {'bread', 'sugar'}
]


# 计算数据集中的频繁项集和关联规则
def apriori(dataset, min_support=0.3, min_confidence=0.6):
    items = {}
    transactions = len(dataset)
    freq_items = {}

    # 计算单个物品的支持度
    for transaction in dataset:
        for item in transaction:
            items[frozenset([item])] = items.get(frozenset([item]), 0) + 1

    # 过滤掉不满足支持度要求的物品
    freq_items[1] = {item: support / transactions for item, support in items.items() if
                     support / transactions >= min_support}

    k = 2
    while len(freq_items[k - 1]) > 0:
        # 生成候选项集
        candidate_items = set([item1.union(item2) for item1 in freq_items[k - 1] for item2 in freq_items[k - 1] if
                               len(item1.union(item2)) == k])

        # 计算候选项集的支持度
        freq_items[k] = {}
        for transaction in dataset:
            for candidate in candidate_items:
                if candidate.issubset(transaction):
                    freq_items[k][candidate] = freq_items[k].get(candidate, 0) + 1

        # 过滤掉不满足支持度要求的项集
        freq_items[k] = {item_set: support / transactions for item_set, support in freq_items[k].items() if
                         support / transactions >= min_support}

        k += 1

    return freq_items


result = apriori(dataset)

# 输出频繁项集
for k, items in result.items():
    print(f"频繁{k}-项集:")
    for item_set, support in items.items():
        print(f"{item_set}:支持度{support}")


# 输出关联规则
def generate_rules(freq_items, min_confidence=0.6):
    rules = []
    for k, items in freq_items.items():
        if k < 2:
            continue
        for item_set, support in items.items():
            for subset in combinations(item_set, k - 1):
                remaining = item_set.difference(subset)
                confidence = support / freq_items[k - 1][frozenset(subset)]
                if confidence >= min_confidence:
                    rules.append((subset, remaining, confidence))
    return rules


rules = generate_rules(result)

print("\n关联规则:")
for rule in rules:
    antecedent, consequent, confidence = rule
    print(f"{antecedent} => {consequent} :置信度{confidence*100:.2f}%")